import time

from PyQt5.QtCore import *
import sys
from enum import Enum
from ai.img_identify import *
import cv2
from robot_arm.my_serial import *
import threading
import socket

###########################################
# 摄像头的编号和串口名需要根据实际情况进行修改
# 串口名
SERIAL_NAME = 'COM1'
# 摄像头编号
VIDEO_NUM = 1
############################################

# 控制机械臂抓取的动作组
CONTROL_ROBOTIC_ARM_GRAB_DATA = "555505060{}0100"

# 创建串口对象
my_serial = MySerial(SERIAL_NAME, baudrate=9600, timeout=1)
# 创建串口接收远程
t_serial = threading.Thread(target=my_serial.receive_msg)
t_serial.start()


class MoveType(Enum):
    MAX_MOVE_TYPE = 0
    MIN_MOVE_TYPE = 1
    FIX_MOVE_TYPE = 2


class HandleThread(QThread):
    photoSignal = pyqtSignal()
    identifySignal = pyqtSignal(int)

    def __init__(self):
        super().__init__()

        # 为了让讲师好讲，UDP的不再封装成类了
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.socket.bind(('0.0.0.0', 8888))
        receive_thread = threading.Thread(target=self.receive_messages)
        receive_thread.daemon = True
        receive_thread.start()
        print('启动UDP接收线程')

        self.val = 0
        # 默认最大搬运
        self.mov_type = MoveType.MAX_MOVE_TYPE
        # 默认搬运2
        self.fix_num = 2
        self.current_state = 0
        try:
            self.vs = cv2.VideoCapture(VIDEO_NUM)
            if not self.vs.isOpened():
                raise Exception("无法打开摄像头")
        except Exception as e:
            print(f"初始化摄像头失败: {e}")
            sys.exit()

        self.image_processing = ImageProcessing()

    # 这个线程函数只是为了对接大模型的UDP服务的
    def receive_messages(self):
        while True:
            try:
                data, address = self.socket.recvfrom(1024)
                print(f"收到来自 {address} 的消息: {data.decode()}")

                # 业务逻辑在这里处理，直接发送55协议控制机械臂运动即可
                move_cmd = 5
                byte_stream = bytes.fromhex(CONTROL_ROBOTIC_ARM_GRAB_DATA.format(move_cmd))
                my_serial.send_msg(byte_stream)

            except Exception as e:
                print(f"接收错误: {e}")
                break

    def set_move_type(self, mov_type, num=2):
        self.mov_type = mov_type
        self.fix_num = num

    def run(self):
        __source_location = 0
        self.current_state = 0

        while True:
            if self.current_state == 0:
                # 先把机械臂移动到仓库1
                my_serial.recv_msg = ""
                byte_stream = bytes.fromhex(CONTROL_ROBOTIC_ARM_GRAB_DATA.format(0))
                my_serial.send_msg(byte_stream)
                print("控制机械臂移动。")
                self.current_state = 1
            elif self.current_state == 1:
                # print('len=', len(my_serial.recv_msg))
                if len(my_serial.recv_msg) >= 4 and my_serial.recv_msg[2] == 0x02 and my_serial.recv_msg[3] == 0x07:
                    print("机械臂已到达仓库1。")
                    self.current_state = 2
                    my_serial.recv_msg = ""
                else:
                    print("等待机械臂到达仓库1")
                    time.sleep(1)
                    continue
            elif self.current_state == 2:
                # input("按下任意键开始拍摄照片...")
                for i in range(30):
                    ret, frame = self.vs.read()
                cv2.imwrite("./pic/result.jpg", frame)
                # cv2.imshow("frame", frame)
                # cv2.waitKey(0)
                print("拍摄照片完成，照片保存在pic目录下result.jpg，开始识别...")
                self.photoSignal.emit()

                # 不接机械臂打开这句测试
                # frame = cv2.imread("./pic/normal.jpg")
                image_thresh, cargo_location = self.image_processing.image_position(frame)
                cargo_location_sort = self.image_processing.image_sort(cargo_location)
                rec_result = self.image_processing.image_recognize(cargo_location_sort, frame)
                if len(rec_result) == 0:
                    print("未识别到货物")
                    time.sleep(1)
                    continue
                print("识别结束，结果为{}".format(rec_result))
                # cv2.imshow('结果', img_end_list[1])
                # cv2.waitKey(0)

                # 接下来根据搬运方式进行标定位置
                if self.mov_type == MoveType.MAX_MOVE_TYPE:
                    print('最大搬运')
                    list_sort = sorted(rec_result.items(), key=lambda kv: (kv[1], kv[0]), reverse=True)
                    print("排序结束，结果为{}".format(list_sort))
                    __source_location = list_sort[0][0] + 1
                    print(f'最大位置为{__source_location}')
                elif self.mov_type == MoveType.MIN_MOVE_TYPE:
                    print('最小搬运')
                    list_sort = sorted(rec_result.items(), key=lambda kv: (kv[1], kv[0]), reverse=False)
                    print("排序结束，结果为{}".format(list_sort))
                    print(list_sort[0][0])
                    __source_location = list_sort[0][0] + 1
                    print(f'最小位置为{__source_location}')
                elif self.mov_type == MoveType.FIX_MOVE_TYPE:
                    print('固定搬运')
                    # 遍历字典
                    for key, value in rec_result.items():
                        if value == self.fix_num:
                            __source_location = key + 1
                            print(f'固定位置为{__source_location}')
                else:
                    print("未知的搬运方式")
                    break

                self.identifySignal.emit(__source_location - 1)
                self.current_state = 3
            elif self.current_state == 3:
                my_serial.recv_msg = ""
                byte_stream = bytes.fromhex(CONTROL_ROBOTIC_ARM_GRAB_DATA.format(__source_location))
                my_serial.send_msg(byte_stream)
                self.current_state = 4
            elif self.current_state == 4:
                if len(my_serial.recv_msg) >= 4 and my_serial.recv_msg[2] == 0x02 and my_serial.recv_msg[3] == 0x07:
                    print("机械臂已搬运完毕。")
                    break
                else:
                    print("等待机械臂搬运完成")
                    time.sleep(1)
                    continue
            else:
                print("未知的状态")
                break
